#include "lowdiscrepancy.h"
#include "../core/sampling.h"

LDSampler::LDSampler(int xstart, int xend, int ystart, int yend, int spp, bool lns, bool ts) :
			Sampler(xstart, xend, ystart, yend, spp, lns, ts) {
	xPos = xstart;
	yPos = ystart;
	if(!IsPowerOf2(spp)) {
		Warning("LDSampler spp roundup to pow 2");
		nPixelSamples = RoundUpPow2(spp);
	} 
	else
		nPixelSamples = spp;

	// Allocate memory for samples

	imageSamples = new float[5*nPixelSamples];
	lensSamples = imageSamples + 2*nPixelSamples;
	timeSamples = lensSamples + 2*nPixelSamples;
	
	memset(timeSamples, 0, sizeof(float)*nPixelSamples);
	memset(lensSamples, 0, sizeof(float)*2*nPixelSamples);
	
	oneDSamples = twoDSamples = NULL;
	// Initialize rng
	rng = new RNG(xstart + ystart * (xend-xstart));
}

LDSampler::~LDSampler() {
	SAFE_DELETE_ARRAY(imageSamples);
	SAFE_DELETE(rng);
	if(oneDSamples) {
		for(int i = 0; i < n1d; i++)
			SAFE_DELETE_ARRAY(oneDSamples[i]);
		SAFE_DELETE_ARRAY(oneDSamples);
	}
	if(twoDSamples) {
		for(int i = 0; i < n2d; i++)
			SAFE_DELETE_ARRAY(twoDSamples[i]);
		SAFE_DELETE_ARRAY(twoDSamples);
	}
}

Sampler *LDSampler::GetSubSampler(int num, int count) {
	int x0, x1, y0, y1;
	ComputeSubWindow(num, count, &x0, &x1, &y0, &y1);
	if (x0 == x1 || y0 == y1) 
		return NULL;
	return new LDSampler(x0, x1, y0, y1, nPixelSamples, needLensSamples, needTimeSamples);
}

int LDSampler::GetMoreSamples(Sample *sample) {

	if(yPos == yPixelEnd)
		return 0;

	if(!oneDSamples) {
		// Allocate memory for oneDSamples and twoDSamples
		Sample &samp = sample[0];
		n1d = samp.n1D.size();
		oneDSamples = new float *[n1d];
		for(int i = 0; i < n1d; i++)
			oneDSamples[i] = new float[samp.n1D[i] * nPixelSamples];
		n2d = samp.n2D.size();
		twoDSamples = new float *[n2d];
		for(int i = 0; i < n2d; i++)
			twoDSamples[i] = new float[2*samp.n2D[i] * nPixelSamples];
	}

	// Generate LD samples

	LDShuffleScrabled2D(1, nPixelSamples, imageSamples, rng);
	if(needLensSamples)
		LDShuffleScrabled2D(1, nPixelSamples, lensSamples, rng);
	if(needTimeSamples)
		LDShuffleScrabled1D(1, nPixelSamples, timeSamples, rng);

	for(int i = 0; i < n1d; i++)
		LDShuffleScrabled1D(sample->n1D[i], nPixelSamples, oneDSamples[i], rng);
	for(int i = 0; i < n2d; i++)
		LDShuffleScrabled2D(sample->n2D[i], nPixelSamples, twoDSamples[i], rng);
	
	for(int i = 0; i < nPixelSamples; i++) {
		sample[i].imageSample = Point2(imageSamples[2*i] + xPos, imageSamples[2*i+1] + yPos);
		sample[i].lensSample = Point2(lensSamples[2*i], lensSamples[2*i+1]);
		sample[i].time = timeSamples[i];

		// Copy integrator samples into _samples[i]_
		for (int j = 0; j < n1d; ++j) {
			int startSamp = sample[i].n1D[j] * i;
			for (int k = 0; k < sample[i].n1D[j]; ++k)
				sample[i].oneD[j][k] = oneDSamples[j][startSamp+k];
		}
		for (int j = 0; j < n2d; ++j) {
			int startSamp = 2 * sample[i].n2D[j] * i;
			for (int k = 0; k < sample[i].n2D[j]; k++)
				sample[i].twoD[j][k] = Point2(twoDSamples[j][startSamp+2*k],twoDSamples[j][startSamp+2*k+1]);
		}

	}
	
	// advance
	xPos++;
	if(xPos == xPixelEnd) {
		xPos = xPixelStart;
		yPos++;
	}

	return nPixelSamples;
}